package agentd_mux1

import (
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"sync"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"wiseverds.com/frame/webport/src/wsmux"
)

type WSConnInfo struct {
	Number  uint32
	ConnID  string
	Service string
	Muxer   *wsmux.Session
}

type TProxy struct {
	m               sync.RWMutex
	onSessionRemove func(string)
	upgrader        websocket.Upgrader
	urlPrefix       string
	map1            map[string]*websocket.Conn
	sp2Client       *wsmux.Session
	sp2ConnChan     chan net.Conn
	clients         map[string]*WSConnInfo
}

type LocalConn struct {
	ID         uint64
	Conn       net.Conn
	ServerName string
	ClientAddr string
}

// Config contains the run time parameters for the proxy
type Config struct {
	// Upgrader is a websocket.Upgrader instance which is used to upgrade incoming
	// websocket connections from Clients.
	Upgrader websocket.Upgrader

	// the prefix for publicly accessible URLs (used to generate the URLs sent
	// to clients)
	URLPrefix string
}

func NewProxy(conf Config) (*TProxy, error) {
	p := &TProxy{
		upgrader:  conf.Upgrader,
		urlPrefix: strings.TrimSuffix(conf.URLPrefix, "/"),
		map1:      make(map[string]*websocket.Conn),
		clients:   make(map[string]*WSConnInfo),
	}

	return p, nil
}

func (p *TProxy) GetWSClientList() []*WSConnInfo {
	list := make([]*WSConnInfo, 0)

	for _, info := range p.clients {
		list = append(list, info)
	}

	return list
}

// setSessionRemoveHandler sets a function which is called when a wsmux Session is removed from
// the proxy due to closure or error.
func (p *TProxy) setSessionRemoveHandler(h func(string)) {
	p.m.Lock()
	defer p.m.Unlock()
	p.onSessionRemove = h
}

// ServeHTTP implements http.Handler so that the proxy may be used as a handler in a Mux or http.Server
func (p *TProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Printf("ServeHTTP: Host=%s Path=%s RemoteAddr=%s", r.Host, r.URL.Path, r.RemoteAddr)

	// try to match viewer requests to <urlPrefix>/<id>/<path> (ignoring
	// the host field)
	uri := r.URL.RequestURI()

	prefix := p.urlPrefix + "/"
	if !strings.HasPrefix(uri, prefix) {
		log.Printf("invalid request url: %s", uri)
		http.NotFound(w, r)
		return
	}

	real_uri := uri[len(prefix):]

	index := strings.Index(real_uri, "/")
	if index < 0 {
		log.Printf("invalid request url: %s", uri)
		http.NotFound(w, r)
		return
	}

	id := real_uri[:index]
	path := real_uri[index:]
	p.serveRequest(w, r, id, path)
}

// serveRequest serves tunnel endpoints to viewers
func (p *TProxy) serveRequest(w http.ResponseWriter, r *http.Request, id string, path string) {
	// log new request arrival
	log.Printf("serveRequest: id=%s host=%s path=%s remoteAddr=%s URL=%s", id, r.Host, path, r.RemoteAddr, r.URL)

	service := r.Header.Get("X-Webport-Service")
	switch service {
	case "sp1":
		fmt.Printf("X-Webport-Service: sp1\n")
		p.serveRequestSp1(w, r, id, path)
	case "":
		fmt.Printf("X-Webport-Service error: (empty)\n")
		// return 404
		http.NotFound(w, r)
		return
	default:
		fmt.Printf("X-Webport-Service error: %s\n", service)
		// return 404
		http.NotFound(w, r)
		return
	}
}

func (p *TProxy) serveRequestSp1(w http.ResponseWriter, r *http.Request, id string, path string) {
	header := make(http.Header)
	url1 := p.urlPrefix + "/" + id
	header.Set("X-Webport-Client-Url", url1)

	connId := p.generateConnectionID()
	header.Set("X-Webport-Connection-Id", connId)

	connNum, err := p.generateConnectionNumber()
	if err != nil {
		log.Printf("failed to generate a connection number: %v", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	conn, err := p.upgrader.Upgrade(w, r, header)
	if err != nil {
		log.Printf("failed to upgrade: %v", err)
		return
	}

	p.map1[connId] = conn

	// generate server config
	confServer := wsmux.Config{
		StreamBufferSize: 4 * 1024,
		CloseCallback: func() {
			if p.onSessionRemove != nil {
				p.onSessionRemove(id)
			}
		},
	}

	server := wsmux.Server(conn, confServer)

	connInfo := &WSConnInfo{
		Number:  connNum,
		ConnID:  connId,
		Service: "sp1",
		Muxer:   server,
	}

	log.Printf("agentd_mux1: connInfo: %v", connInfo)

	p.clients[connId] = connInfo

	// accept new streams from this channel
	strChan := make(chan net.Conn, 1)

	go func() {
		for {
			// This simulates http.Server's Serve method
			log.Printf("try to accept a new connection ...")
			stream, err := server.Accept()
			if err != nil {
				log.Printf("server.Accept error: %v", err)
				// _, ok := err.(net.Error)
				// if !ok {
				// 	log.Printf("server.Accept error: %v", err)
				// }
				return
			}

			strChan <- stream
		}
	}()

	defer func() {
		_ = server.Close()
		close(strChan)
	}()

	for {
		select {
		case stream := <-strChan:
			log.Printf("Accepted a new connection")
			f := &forwarder{
				stream: stream,
			}

			go f.forward()
		}
	}
}

func (p *TProxy) generateConnectionNumber() (uint32, error) {
	for i := 0; i < 1000; i += 1 {
		num := rand.Uint32()

		found := false
		for _, info := range p.clients {
			if info.Number == num {
				found = true
				break
			}
		}

		if !found {
			return num, nil
		}
	}

	return 0, errors.New("failed to generate a connection number")
}

// Generate a new connection ID.
// This is used to generate the URL for the viewer to connect to.
func (p *TProxy) generateConnectionID() string {
	return uuid.NewString()
}

//----------------------------------------------------------------------------------------------
// agentd_ss1 (socks5) handler
//----------------------------------------------------------------------------------------------

func (p *TProxy) Resolve(name string) (net.IP, error) {
	info, ok := p.clients[name]
	if !ok {
		log.Printf("agentd_mux1: Resolve: unknown host: %s", name)
		return nil, errors.New("invalid host")
	}

	num := info.Number
	ip := []byte{0, 0, 0, 0}
	ip[0] = byte(num >> 24)
	ip[1] = byte((num >> 16) & 0xff)
	ip[2] = byte((num >> 8) & 0xff)
	ip[3] = byte(num & 0xff)
	return ip, nil

	// ip := "127.0.0.1"
	// addr := &net.IPAddr{IP: net.ParseIP(ip)}
	// return ctx, addr.IP, nil

	// ip := []byte{127, 0, 0, 1}
	// return ip, nil
}

func parseAddr(addr string) (net.IP, int, error) {
	regexp1 := regexp.MustCompile(`^(\d+)\.(\d+)\.(\d+)\.(\d+):(\d+)$`)
	matches := regexp1.FindStringSubmatch(addr)
	a, _ := strconv.Atoi(matches[1])
	b, _ := strconv.Atoi(matches[2])
	c, _ := strconv.Atoi(matches[3])
	d, _ := strconv.Atoi(matches[4])
	ip := []byte{byte(a), byte(b), byte(c), byte(d)}
	port, _ := strconv.Atoi(matches[5])
	return ip, port, nil
}

func (p *TProxy) Dial(addr string) (net.Conn, error) {
	log.Println("webport_sp1: dial:", addr)

	ip, port, err := parseAddr(addr)
	log.Println("webport_sp1: parseAddr:", ip, port)

	connNum := 0
	connNum += (int(ip[0]) << 24)
	connNum += (int(ip[1]) << 16)
	connNum += (int(ip[2]) << 8)
	connNum += int(ip[3])

	var connInfo *WSConnInfo
	for _, info := range p.clients {
		if info.Number == uint32(connNum) {
			connInfo = info
			break
		}
	}

	if connInfo == nil {
		log.Printf("agentd_mux1: handleConn: unknown host: %s", addr)
		return nil, errors.New("invalid host")
	} else {
		log.Printf("agentd_mux1: handleConn: found host: %v", connInfo)
	}

	stream, err := connInfo.Muxer.Open()
	if err != nil {
		log.Printf("webport_sp1: handleConn: failed to open stream: %v", err)
		return nil, errors.New("failed to open stream")
	}

	err = binary.Write(stream, binary.BigEndian, uint32(port))
	if err != nil {
		return nil, err
	}

	return stream, nil
}
